/*
 * Copyright (c) Kumo Inc. and affiliates.
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <melon/timeout_queue.h>

#include <algorithm>
#include <limits>
#include <vector>

namespace melon {
    TimeoutQueue::Id TimeoutQueue::add(
        int64_t now, int64_t delay, Callback callback) {
        Id id = nextId_++;
        timeouts_.insert({id, now + delay, -1, std::move(callback)});
        return id;
    }

    TimeoutQueue::Id TimeoutQueue::addRepeating(
        int64_t now, int64_t interval, Callback callback) {
        Id id = nextId_++;
        timeouts_.insert({id, now + interval, interval, std::move(callback)});
        return id;
    }

    int64_t TimeoutQueue::nextExpiration() const {
        return (
            timeouts_.empty()
                ? std::numeric_limits<int64_t>::max()
                : timeouts_.get<BY_EXPIRATION>().begin()->expiration);
    }

    bool TimeoutQueue::erase(Id id) {
        return timeouts_.get<BY_ID>().erase(id);
    }

    int64_t TimeoutQueue::runInternal(int64_t now, bool onceOnly) {
        auto &byExpiration = timeouts_.get<BY_EXPIRATION>();
        int64_t nextExp;
        do {
            const auto end = byExpiration.upper_bound(now);
            std::vector<Event> expired;
            std::move(byExpiration.begin(), end, std::back_inserter(expired));
            byExpiration.erase(byExpiration.begin(), end);
            for (const auto &event: expired) {
                // Reinsert if repeating, do this before executing callbacks
                // so the callbacks have a chance to call erase
                if (event.repeatInterval >= 0) {
                    timeouts_.insert(
                        {
                            event.id,
                            now + event.repeatInterval,
                            event.repeatInterval,
                            event.callback
                        });
                }
            }

            // Call callbacks
            for (const auto &event: expired) {
                event.callback(event.id, now);
            }
            nextExp = nextExpiration();
        } while (!onceOnly && nextExp <= now);
        return nextExp;
    }
} // namespace melon
